{-# LANGUAGE ScopedTypeVariables, GADTs, StandaloneDeriving, Rank2Types #-}
{-# LANGUAGE MultiParamTypeClasses, FunctionalDependencies, FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts, GeneralizedNewtypeDeriving #-}
module Yices
    ( Name, Ident
    , Term, Prop, Clause
    , Type
    , Alias, Opaque, Transparent, alias, aconv, unalias
    , Nat
    , Scalar, Refl (refl), Named (sname)
    , stdRefl
    , (!)
    , idt, int, bool, nat, lit
    , ident, ident1, ident2, ident3
    , assert, define, deftype, typeAlias
    , and, or, not, xor, xor'
    , (/\), (\/), non
    , conj, disj, ifte, ulog, ilog, bits
    , (==>), (<==>)
    , ge', ge, le', le, plus', plus, minus', minus
    , gt', gt, lt', lt
    , mul', mul, eq', eq, iff', iff
    , printScript, printBddSat, printBddReach
    ) where

import Prelude hiding (and, or, not)
import Text.PrettyPrint.HughesPJ hiding (int)
import Data.Char
import Data.Ord
import Data.List hiding (and, or)
import qualified Data.Set as S

type Name = String
type Prop = Term Bool
data Ident t = Ident Name
data Type = Scalar Name | Type :-> Type

data NatTag
type Nat = Alias Opaque NatTag Int

data Opaque
data Transparent
newtype Alias op tag t = Alias { unalias :: t } deriving (Bounded,Num,Enum,Integral,Eq,Ord,Real,Scalar)

instance (Show t) => Show (Alias op tag t) where
    show (Alias x) = show x

data Term t where
    Lit :: Scalar t => t -> Term t
    Var :: Refl t => Ident t -> Term t
    App :: (Refl a, Refl b) => Term (a -> b) -> Term a -> Term b

data Clause where
    TypeAlias :: (Named (Alias Opaque tag t), Named t) => Alias Opaque tag t -> Clause
    Deftype :: (Scalar a, Named a) => a -> Clause
    Definition :: Refl t => Ident t -> Clause
    Assertion :: Prop -> Clause

instance Show Clause where
    show = render . pretty . simpl . code

data SExp = S String [SExp] deriving Show

class Named a where
    sname :: a -> String

class (Bounded a, Enum a, Show a) => Scalar a where
    srepr :: a -> String
    srepr = map toLower . show

    senum :: [a]
    senum = [minBound..maxBound]

instance Named Bool where   sname _ = "bool"
instance Scalar Bool where  senum = error "cannot declare bool"

instance Named Int where    sname _ = "int"
instance Scalar Int where   senum = error "cannot declare int"

instance Named (Alias Opaque NatTag Int) where    sname _ = "nat"

instance Named t => Named (Alias Transparent tag t) where
    sname (Alias x) = sname x

class Refl a where
    refl :: a -> Type

instance Refl Bool where refl = stdRefl
instance Refl Int where refl = stdRefl
instance Refl (Alias Opaque NatTag Int) where refl = stdRefl
instance Refl t => Refl (Alias Transparent tag t) where refl (Alias x) = refl x

instance (Refl a, Refl b) => Refl (a -> b) where
    refl _ = (refl (undefined :: a)) :-> (refl (undefined :: b))

stdRefl :: (Named a, Scalar a) => a -> Type
stdRefl = Scalar . sname

slit :: String -> SExp
slit x = S x []

class Code a where
    code :: a -> SExp

instance Code Type where
    code (Scalar x) = slit x 
    code (p :-> q) = reassoc $ S "->" [code p, code q]
      where
        reassoc :: SExp -> SExp
        reassoc (S "->" []) = error "reassoc: Fishy type, arrow not applied."
        reassoc (S "->" xs) =
            case (init xs, last xs) of
                (ys, S "->" zs) -> reassoc $ S "->" (ys ++ zs)
                _              -> S "->" xs
        reassoc x = x

instance Code Clause where
    code (Assertion t) = S "assert" [code t]
    code (TypeAlias x) = S "define-type" $ map slit
        [sname x, sname (inner x)]
      where
        inner :: Alias Opaque tag t -> t
        inner _ = undefined
    code (Deftype x) = S "define-type"
        [slit $ sname x, S "scalar" $ map (slit . srepr . idAs x) senum]
      where
        idAs :: t -> t -> t
        idAs _ = id
    code (Definition i@(Ident n)) = S "define"
        [slit n, slit "::", code (refl $ typeOf i)]
      where
        typeOf :: Ident t -> t
        typeOf _ = undefined

instance Refl a => Code (Term a) where
    code (Lit x) = slit $ srepr x
    code (Var (Ident n)) = slit n
    code (App p q) = let S n xs = code p in S n (xs ++ [code q])

infixl 4 !
class Applicable p q r | p q -> r where
    (!) :: p -> q -> r

instance (Refl a, Refl b) => Applicable (Term (a -> b)) (Term a) (Term b)
    where (!) = App
instance (Refl a, Refl b) => Applicable (Ident (a -> b)) (Term a) (Term b)
    where p ! q = App (idt p) q
instance (Refl a, Refl b) => Applicable (Term (a -> b)) (Ident a) (Term b)
    where p ! q = App p (idt q)
instance (Refl a, Refl b) => Applicable (Ident (a -> b)) (Ident a) (Term b)
    where p ! q = App (idt p) (idt q)

assert :: Term Bool -> Clause
assert = Assertion

define :: Refl t => Ident t -> Clause
define = Definition

deftype :: (Named a, Scalar a) => a -> Clause
deftype = Deftype

typeAlias :: (Named (Alias Opaque tag t), Named t) => Alias Opaque tag t -> Clause
typeAlias = TypeAlias

idt :: (Refl t) => Ident t -> Term t
idt = Var

int :: Int -> Term Int
int = Lit

nat :: Nat -> Term Nat
nat = Lit

bool :: Bool -> Term Bool
bool = Lit

ident :: Refl t => Name -> Ident t
ident = Ident

lit :: Scalar a => a -> Term a
lit = Lit

alias :: a -> Alias op tag a
alias = Alias

aconv :: Alias op1 tag1 a -> Alias op2 tag2 a
aconv = alias . unalias

unlift1 :: (Refl a, Refl b, Scalar a)
    => Ident (a -> b) -> a -> Term b
unlift1 i x = i ! lit x

ident1 :: (Refl a, Refl b, Scalar a)
    => Name -> (Ident (a -> b), a -> Term b)
ident1 n = (ident n, unlift1 $ ident n)

unlift2 :: (Refl a, Refl b, Refl c, Scalar a, Scalar b)
    => Ident (a -> b -> c) -> a -> b -> Term c
unlift2 i x y = i ! lit x ! lit y

ident2 :: (Refl a, Refl b, Refl c, Scalar a, Scalar b)
    => Name -> (Ident (a -> b -> c), a -> b -> Term c)
ident2 n = (ident n, unlift2 $ ident n)

unlift3 :: (Refl a, Refl b, Refl c, Refl d, Scalar a, Scalar b, Scalar c)
    => Ident (a -> b -> c -> d) -> a -> b -> c -> Term d
unlift3 i x y z = i ! lit x ! lit y ! lit z

ident3 :: (Refl a, Refl b, Refl c, Refl d, Scalar a, Scalar b, Scalar c)
    => Name -> (Ident (a -> b -> c -> d), a -> b -> c -> Term d)
ident3 n = (ident n, unlift3 $ ident n)

pretty :: SExp -> Doc
pretty (S n []) = text n
--pretty (S x xs@(_:_:_:_:_)) = parens $ hsep [text x, vcat $ map pretty xs]
pretty (S n xs) = parens . hsep $ map pretty (slit n:xs)

simpl :: SExp -> SExp
simpl (S "not" [S "not" [sub]]) = simpl sub 
simpl (S op1 (S op2 xs : ys))
    | (op1 == op2) && assoc op1 = simpl $ S op1 (xs ++ ys)
simpl (S n xs) = S n (map simpl xs)

assoc :: String -> Bool
assoc x = x `elem` words "+ * and or"

-- *** General utilities ***

conj :: [Prop] -> Prop
conj [] = bool True
conj xs = foldl1 (/\) xs

disj :: [Prop] -> Prop
disj [] = bool False
disj xs = foldl1 (\/) xs

infixr 2 ==>
impl' :: Ident (Bool -> Bool -> Bool)
(==>) :: Prop -> Prop -> Prop
impl' = ident "implies"
p ==> q = idt impl' ! p ! q

equiv :: Ident (Bool -> Bool -> Bool)
equiv = ident "="

infixr 2 <==>
(<==>) :: Prop -> Prop -> Prop
f <==> g = equiv ! f ! g

ulog :: Int -> Int
ulog = length . takeWhile (>1) . iterate (`div` 2)

ilog :: Int -> Int
ilog = (+1) . ulog . (subtract 1)

bits :: Int -> (Int -> Prop) -> Prop
bits n c = conj $ zipWith id (binary n) (map c [0..])

binary :: Int -> [Prop -> Prop]
binary 0 = [non]
binary 1 = [id]
binary n
    | n `mod` 2 == 0 = non : binary (n `div` 2)
    | otherwise     = id  : binary (n `div` 2)

not :: Ident (Bool -> Bool)
not = ident "not"

and :: Ident (Bool -> Bool -> Bool)
and = ident "and"

or :: Ident (Bool -> Bool -> Bool)
or = ident "or"

infixl 3 /\
(/\) :: Term Bool -> Term Bool -> Term Bool
p /\ q = idt and ! p ! q

infixl 3 \/
(\/) :: Term Bool -> Term Bool -> Term Bool
p \/ q = idt or ! p ! q

non :: Term Bool -> Term Bool
non = (idt not !)

ge' :: Ident (Nat -> Nat -> Bool)
(ge',ge) = ident2 ">="

le' :: Ident (Nat -> Nat -> Bool)
(le',le) = ident2 "<="

gt' :: Ident (Nat -> Nat -> Bool)
(gt',gt) = ident2 ">"

lt' :: Ident (Nat -> Nat -> Bool)
(lt',lt) = ident2 "<"

eq' :: Ident (Nat -> Nat -> Bool)
(eq',eq) = ident2 "="

plus' :: Ident (Nat -> Nat -> Nat)
(plus',plus) = ident2 "+"

mul' :: Ident (Nat -> Nat -> Nat)
(mul',mul) = ident2 "*"

minus' :: Ident (Nat -> Nat -> Nat)
(minus', minus) = ident2 "-"

iff' :: Ident (Bool -> Bool -> Bool)
(iff',iff) = ident2 "iff"

xor' :: Ident (Bool -> Bool -> Bool)
(xor',xor) = ident2 "xor"

ifte :: Prop -> Prop -> Prop -> Prop
ifte p t e = idt ifte' ! p ! t ! e
  where
    ifte' :: Ident (Bool -> Bool -> Bool -> Bool)
    ifte' = ident "if"

printScript :: [Clause] -> IO ()
printScript = putStr . unlines . (++ ["(check)"]) . map show

printBddReach :: Prop -> Prop -> Prop -> IO ()
printBddReach initial trans final = putStr . unlines $
    [ "(benchmark bddreach"
    , "  :logic REACH"
    , "  :variables (VARS)"
    , "  :initial_state " ++ format initial
    , "  :transition_relation " ++ format trans
    , "  :final_state " ++ format final
    , ")"
    ]
  where
    format = render . pretty . squash . simpl . code

printBddSat :: Prop -> IO ()
printBddSat p = putStr . unlines $
    [ "(benchmark bddsat"
    , "  :logic SAT"
    , "  :variables (" ++ unwords vars ++ ")"
    , "  :formula " ++ formula
    , ")"
    ]
  where
    formula = render . pretty $ sexp
    sexp = squash . simpl . code $ p
    vars = reorder $ S.elems (variables sexp S.\\ S.fromList bddOps)

reorder :: [String] -> [String]
reorder = sortBy (comparing trans)
  where
    spaced = map (\x -> if (x == '_') then ' ' else x)
    trans :: String -> (String, [Int])
    trans x = let (f:as) = words $ spaced x in (f,map read as)

bddOps :: [String]
bddOps = words "and or iff not xor implies"

squash :: SExp -> SExp
squash (S op xs)
    | op `elem` bddOps = S op (map squash xs)
    | otherwise = slit $ intercalate "_" (op : map unop xs)
  where
    unop :: SExp -> String
    unop (S x []) = x
    unop (S x xs) = error "unop: complex value for bdd solver" -- "(" ++ intercalate "_" (x : map unop xs) ++ ")"

variables :: SExp -> S.Set String
variables (S x xs) = S.insert x $ S.unions (map variables xs)
